from corre_func_v2 import * 
import constant as ct
# VdV_PxPyPz = [( 0,0,0 ),(1,0,0),(0,1,0),(0,0,1),
#         (1,1,0),(1,0,1),(0,1,1),(-1,1,0),(-1,0,1),(0,-1,1),
#         (1,1,1),(-1,1,1),(1,-1,1),(1,1,-1), # len=13
#         (-1,0,0),(0,-1,0),(0,0,-1),
#         (-1,-1,0),(-1,0,-1),(0,-1,-1),(1,-1,0),(1,0,-1),(0,1,-1),
#         (-1,-1,-1),(1,-1,-1),(-1,1,-1),(-1,-1,1)]  # len=27
VdV_PxPyPz = [( 0,0,0 ),(1,0,0),(0,1,0),(0,0,1),
        (1,1,0),(1,0,1),(0,1,1),(-1,1,0),(-1,0,1),(0,-1,1),
        (-1,0,0),(0,-1,0),(0,0,-1),
        (-1,-1,0),(-1,0,-1),(0,-1,-1),(1,-1,0),(1,0,-1),(0,1,-1)
        ]  # len=27
VdV_idx={}
for i, p in enumerate(VdV_PxPyPz):
    VdV_idx[p]=i  
class corr_fun: 
    def __init__(self, mom, dire):
        self.mom=mom 
        self.dir=dire 
        self.new_file={}
        self.corr_th=th.zeros((len(mom),ct.Nt,ct.Nt),dtype=th.complex128,device=device)
    def save(self):
        # print(self.new_file)
        corr_np = self.corr_th.to("cpu").numpy()
        for key in self.new_file.keys():
            np.savetxt(self.new_file[key],corr_np[key].view(float).reshape(-1,2))

class corr_fun_mul: 
    def __init__(self, mom, dire):
        self.mom=mom 
        self.dir=dire 
        self.new_file={}
    def save(self):
        # print(self.new_file)
        corr_np = self.corr_th.to("cpu").numpy()
        for key in self.new_file.keys():
            np.savetxt(self.new_file[key], corr_np[key].view(float).reshape(-1,2))

class corr_fun_pi(corr_fun):
    def __init__(self,mom,dire):
        corr_fun.__init__(self, mom,dire)
    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for t_se in range(ct.Nt):
            for i, mom in enumerate(self.mom):
                P_se,P_sk = mom;
                Pxi,Pyi,Pzi = P_se
                Pxf,Pyf,Pzf = P_sk
                filen= "%s/corr_Pi_Px%sPy%sPz%s_Px%sPy%sPz%s.conf%s.dat"%(self.dir, 
                            Pxf,Pyf,Pzf,Pxi,Pyi,Pzi,conf_id)
                if( (not os.path.exists(filen))):
                    self.new_file[i]=(filen) 
                    self.corr_th[i,t_sk,t_se] = Pi_PfPi(perams_u_th[t_sk,t_se],
                            VdV_th[VdV_idx[P_sk],t_sk],VdV_th[VdV_idx[P_se],t_se])


class corr_fun_k(corr_fun):
    def __init__(self,mom,dire):
        corr_fun.__init__(self,  mom,dire)
    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for t_se in range(ct.Nt):
            for i, mom in enumerate(self.mom):
                P_se,P_sk = mom;
                Pxi,Pyi,Pzi = P_se
                Pxf,Pyf,Pzf = P_sk
                filen= "%s/corr_K_Px%sPy%sPz%s_Px%sPy%sPz%s.conf%s.dat"%(self.dir, 
                            Pxf,Pyf,Pzf,Pxi,Pyi,Pzi,conf_id)
                if( (not os.path.exists(filen))):
                    self.new_file[i]=(filen) 
                    self.corr_th[i,t_sk,t_se] = K_PfPi(perams_u_th[t_sk,t_se],
                            perams_s_th[t_se],
                            VdV_th[VdV_idx[P_sk],t_sk],VdV_th[VdV_idx[P_se],t_se])

class corr_fun_ki_mul(corr_fun_mul):
    def __init__(self,mom,dire,gamma,conf_id,VdV_th):
        corr_fun.__init__(self,  mom,dire)
        self.gamma=gamma
        self.corr_th=th.zeros((len(self.mom),len(self.mom),len(self.gamma),
                              len(self.gamma),ct.Nt ),dtype=th.complex128,
                              device=device)
        for i, momsk in enumerate(self.mom):
            for j in range(i,len(self.mom)):
                for n,ga1 in enumerate( self.gamma):
                    for m,ga2 in enumerate( self.gamma):
                        filen= f"{self.dir}/corr_K"+ga1+ga2+\
                            "_Px%sPy%sPz%s_Px%sPy%sPz%s.conf%s.dat"%(
                            *momsk,*(self.mom[j]),conf_id)
                        self.new_file[(i,j,n,m)]=filen
        self.Ga = th.zeros((4,4,len(self.gamma)), dtype=th.complex128,
                            device=device)
        for k,ga1 in enumerate( self.gamma):
            self.Ga[:,:,k] = Gammas[ga1]
        self.VdV_Th_se={}
        for i, _ in enumerate(self.mom):
            self.VdV_Th_se[i]=th.zeros((ct.Nt,ct.Nev1,ct.Nev1,len(self.mom)-i)
                                       ,dtype=th.complex128, device=device)
            for j, momse in enumerate(self.mom[i:]):
                self.VdV_Th_se[i][:,:,:,j] = VdV_th[VdV_idx[momse]]
    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for i, momsk in enumerate(self.mom):
            VdV_th_sk=VdV_th[VdV_idx[momsk]]
            for t_se in range(ct.Nt):
                self.corr_th[i,i:,:,:,(t_sk -t_se)%ct.Nt] += \
                    K_g1g2_2P_mul(perams_u_th[t_sk,t_se], perams_s_th[t_se],
                            VdV_th_sk[t_sk],self.VdV_Th_se[i][t_se],
                            self.Ga,self.Ga)
                

class corr_fun_ki(corr_fun):
    def __init__(self,mom,dire,gamma,conf_id):
        corr_fun.__init__(self,  mom,dire)
        self.gamma=gamma
        self.corr_th=th.zeros((len(mom),ct.Nt),dtype=th.complex128,device=device)
        for i, mom in enumerate(self.mom):
            P_se,P_sk = mom;
            Pxi,Pyi,Pzi = P_se
            Pxf,Pyf,Pzf = P_sk
            filen= f"{self.dir}/corr_K"+self.gamma[0]+self.gamma[1]+"_Px%sPy%sPz%s_Px%sPy%sPz%s.conf%s.dat"%(
                        Pxf,Pyf,Pzf,Pxi,Pyi,Pzi,conf_id)
            self.new_file[i]=(filen) 
    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for t_se in range(ct.Nt):
            for i, mom in enumerate(self.mom):
                P_se,P_sk = mom;
                Pxi,Pyi,Pzi = P_se
                Pxf,Pyf,Pzf = P_sk
                filen= f"{self.dir}/corr_K"+self.gamma[0]+self.gamma[1]+"_Px%sPy%sPz%s_Px%sPy%sPz%s.conf%s.dat"%(
                            Pxf,Pyf,Pzf,Pxi,Pyi,Pzi,conf_id)
                self.corr_th[i,(-t_sk+t_se)%ct.Nt] += K_g1g2_2P(perams_u_th[t_sk,t_se], perams_s_th[t_se],
                            VdV_th[VdV_idx[P_sk],t_sk],VdV_th[VdV_idx[P_se],t_se],self.gamma)
class corr_fun_pik_mul_gi(corr_fun_mul):
    def __init__(self,momsk,momse,dire,conf_id,VdV_th,gammas):
        corr_fun.__init__(self,  momse,dire)
        self.corr_th=th.zeros((3,len(gammas),len(gammas),len(momse),ct.Nt),dtype=th.complex128,device=device)
        self.VdV_Th_se1={}
        self.VdV_Th_se2={}
        self.momse, self.momsk=momse,momsk
        self.gammas=gammas
        self.Gammas = th.zeros((4,4,len(gammas)), dtype=th.complex128,device=device)
        self.Pi_sk,self.K_sk =  momsk
        self.VdV_Th_se1=th.zeros((ct.Nt,ct.Nev1,ct.Nev1,
                                  len(self.momse) ),dtype=th.complex128,
                               device=device)
        self.VdV_Th_se2=th.zeros((ct.Nt,ct.Nev1,ct.Nev1,
                                  len(self.momse) ),dtype=th.complex128,
                               device=device)
        for i, ga in enumerate(self.gammas):
            self.Gammas[:,:,i] = Gammas[ga]
        for i,mom in enumerate(self.momse):
            Pi_se,K_se=mom
            self.VdV_Th_se1[:,:,:,i]= VdV_th[VdV_idx[Pi_se]]
            self.VdV_Th_se2[:,:,:,i]= VdV_th[VdV_idx[K_se]]
            for j,ga1 in enumerate(self.gammas):
                for k,ga2 in enumerate(self.gammas):
                    filenX="%s/PiKX/corr_PiKX_%s%s_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                                ga1,ga2,*self.Pi_sk,*self.K_sk,*Pi_se,*K_se ,conf_id)
                
                    filenH="%s/PiKH/corr_PiKH_%s%s_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                               ga1,ga2, *self.Pi_sk,*self.K_sk,*Pi_se,*K_se ,conf_id)
                    filenA="%s/PiKA/corr_PiKA_%s%s_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                                ga1,ga2,*self.Pi_sk,*self.K_sk,*Pi_se,*K_se ,conf_id)
                    self.new_file[(0,j,k,i)] = filenX 
                    self.new_file[(1,j,k,i)] = filenH 
                    self.new_file[(2,j,k,i)] = filenA 
    def contract(self,perams_u_th,perams_s_th, VdV_th,conf_id,t_sk):
        VdV_th_sk3 = VdV_th[VdV_idx[self.Pi_sk]]
        VdV_th_sk4 = VdV_th[VdV_idx[self.K_sk]]
        for t_se in range(ct.Nt):
            self.corr_th[2,:,:,:,(t_sk -t_se)%ct.Nt] += PiKA_4P_mul_gi(
                        perams_u_th[t_sk,t_se], perams_s_th[t_se],
                        VdV_th_sk3[t_sk],VdV_th_sk4[t_sk],self.VdV_Th_se1[t_se],self.VdV_Th_se2[t_se],self.Gammas)

            self.corr_th[0,:,:,:,(t_sk -t_se)%ct.Nt] += PiKX_4P_mul_gi(
                        perams_u_th[t_sk,t_se], perams_s_th[t_se],
                        VdV_th_sk3[t_sk],VdV_th_sk4[t_sk],self.VdV_Th_se1[t_se],self.VdV_Th_se2[t_se],self.Gammas)

            self.corr_th[1,:,:,:,(t_sk -t_se)%ct.Nt] += PiKH_4P_mul_gi(
                        perams_u_th[t_se,t_se], perams_u_th[t_sk,t_sk],
                        perams_u_th[t_sk,t_se],perams_s_th[t_se],
                        VdV_th_sk3[t_sk],VdV_th_sk4[t_sk],self.VdV_Th_se1[t_se],self.VdV_Th_se2[t_se],self.Gammas)
                    
class corr_fun_pik_mul(corr_fun_mul):
    def __init__(self,mom,dire,conf_id,VdV_th):
        corr_fun.__init__(self,  mom,dire)
        self.corr_th_pikA=th.zeros((len(mom),len(mom),ct.Nt),dtype=th.complex128,device=device)
        self.corr_th_pikX=th.zeros((len(mom),len(mom),ct.Nt),dtype=th.complex128,device=device)
        self.corr_th_pikH=th.zeros((len(mom),len(mom),ct.Nt),dtype=th.complex128,device=device)
        self.corr_th=th.zeros((3,len(mom),len(mom),ct.Nt),dtype=th.complex128,device=device)
        self.VdV_Th_se1={}
        self.VdV_Th_se2={}
        for i,mom in enumerate(self.mom):
            Pi_sk,K_sk =  mom
            self.VdV_Th_se1[i]= th.zeros((ct.Nt,ct.Nev1,ct.Nev1,
                                  len(self.mom) -i),dtype=th.complex128,
                               device=device)
            self.VdV_Th_se2[i]= th.zeros((ct.Nt,ct.Nev1,ct.Nev1,
                                  len(self.mom) -i),dtype=th.complex128,
                               device=device)
            for j in range(i, len(self.mom) ):
                Pi_se,K_se =  self.mom[j]
                self.VdV_Th_se1[i][:,:,:,j-i] = VdV_th[VdV_idx[Pi_se]]
                self.VdV_Th_se2[i][:,:,:,j-i] = VdV_th[VdV_idx[K_se]]
                filenX="%s/PiKX/corr_PiKX_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                            *Pi_sk,*K_sk,*Pi_se,*K_se ,conf_id)
            
                filenH="%s/PiKH/corr_PiKH_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                            *Pi_sk,*K_sk,*Pi_se,*K_se ,conf_id)
                filenA="%s/PiKA/corr_PiKA_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                            *Pi_sk,*K_sk,*Pi_se,*K_se ,conf_id)
                self.new_file[(0,i,j)] = filenX 
                self.new_file[(1,i,j)] = filenH 
                self.new_file[(2,i,j)] = filenA 
    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for i,mom in enumerate(self.mom):
            Pi_sk,K_sk =  mom
            VdV_th_sk3=VdV_th[VdV_idx[Pi_sk]]
            VdV_th_sk4=VdV_th[VdV_idx[K_sk]]
            for t_se in range(ct.Nt):
                self.corr_th_pikA[i,i:,(t_sk -t_se)%ct.Nt] += PiKA_4P_mul(
                            perams_u_th[t_sk,t_se], perams_s_th[t_se],
                            VdV_th_sk3[t_sk],VdV_th_sk4[t_sk],self.VdV_Th_se1[i][t_se],self.VdV_Th_se2[i][t_se])

                self.corr_th_pikX[i,i:,(t_sk -t_se)%ct.Nt] += PiKX_4P_mul(
                            perams_u_th[t_sk,t_se], perams_s_th[t_se],
                            VdV_th_sk3[t_sk],VdV_th_sk4[t_sk],self.VdV_Th_se1[i][t_se],self.VdV_Th_se2[i][t_se])

                self.corr_th_pikH[i,i:,(t_sk -t_se)%ct.Nt] += PiKH_4P_mul(
                            perams_u_th[t_se,t_se], perams_u_th[t_sk,t_sk],
                            perams_u_th[t_sk,t_se],perams_s_th[t_se],
                            VdV_th_sk3[t_sk],VdV_th_sk4[t_sk],self.VdV_Th_se1[i][t_se],self.VdV_Th_se2[i][t_se])
        self.corr_th[0] =self.corr_th_pikX
        self.corr_th[1] =self.corr_th_pikH
        self.corr_th[2] =self.corr_th_pikA
                    
                        
class corr_fun_pikx(corr_fun):
    def __init__(self,mom,dire,conf_id):
        corr_fun.__init__(self,  mom,dire)
        self.corr_th=th.zeros((len(mom),1,ct.Nt),dtype=th.complex128,device=device)
        for i,mom in enumerate(self.mom):
            Pi_se,K_se =  mom[0]
            Pi_sk,K_sk  =  mom[1]
            Pi_Pxi,Pi_Pyi,Pi_Pzi = Pi_se
            Pi_Pxf,Pi_Pyf,Pi_Pzf = Pi_sk
            K_Pxi,K_Pyi,K_Pzi = K_se
            K_Pxf,K_Pyf,K_Pzf = K_sk
            filen= "%s/corr_PiKX_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                        Pi_Pxf,Pi_Pyf,Pi_Pzf,K_Pxf,K_Pyf,K_Pzf,
                        Pi_Pxi,Pi_Pyi,Pi_Pzi,K_Pxi,K_Pyi,K_Pzi,conf_id)
            self.new_file[i]=(filen) 
    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for t_se in range(ct.Nt):
            for i,mom in enumerate(self.mom):
                Pi_se,K_se =  mom[0]
                Pi_sk,K_sk  =  mom[1]
                self.corr_th[i,0,(t_sk -t_se)%ct.Nt] += PiKX_4P(perams_u_th[t_sk,t_se],
                     perams_s_th[t_se],VdV_th[VdV_idx[Pi_sk],t_sk],VdV_th[VdV_idx[K_sk],t_sk],
                                VdV_th[VdV_idx[Pi_se],t_se],VdV_th[VdV_idx[K_se],t_se])
class corr_fun_pikh(corr_fun):
    def __init__(self,mom,dire,conf_id):
        corr_fun.__init__(self,  mom,dire)
        self.corr_th=th.zeros((len(mom),1,ct.Nt),dtype=th.complex128,device=device)
        for i,mom in enumerate(self.mom):
            Pi_se,K_se =  mom[0]
            Pi_sk,K_sk  =  mom[1]
            Pi_Pxi,Pi_Pyi,Pi_Pzi = Pi_se
            Pi_Pxf,Pi_Pyf,Pi_Pzf = Pi_sk
            K_Pxi,K_Pyi,K_Pzi = K_se
            K_Pxf,K_Pyf,K_Pzf = K_sk
            filen= "%s/corr_PiKH_Pi%s%s%sK%s%s%s_Pi%s%s%sK%s%s%s.conf%s.dat"%(self.dir,
                        Pi_Pxf,Pi_Pyf,Pi_Pzf,K_Pxf,K_Pyf,K_Pzf,
                        Pi_Pxi,Pi_Pyi,Pi_Pzi,K_Pxi,K_Pyi,K_Pzi,conf_id)
            self.new_file[i]=(filen) 
    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for t_se in range(ct.Nt):
            for i,mom in enumerate(self.mom):
                Pi_se,K_se =  mom[0]
                Pi_sk,K_sk  =  mom[1]
                self.corr_th[i,0,(t_sk -t_se)%ct.Nt] += PiKH_4P(perams_u_th[t_se,t_se],
                         perams_u_th[t_sk,t_sk],perams_u_th[t_sk,t_se],perams_s_th[t_se],
                         VdV_th[VdV_idx[Pi_sk],t_sk],VdV_th[VdV_idx[K_sk],t_sk],
                                    VdV_th[VdV_idx[Pi_se],t_se],VdV_th[VdV_idx[K_se],t_se])  
class corr_fun_pikkg(corr_fun):
    def __init__(self,mom,dire,gammas,conf_id):
        corr_fun.__init__(self,mom,dire)
        self.gamms=gammas
        for i,mom in enumerate(self.mom):
            Pi,K =  mom[0]
            Kg  =  mom[1]
            filen= f"{self.dir}/corr_PiK_K"+self.gamms[-1]+"_Pi%s%s%sK%s%s%s_Kg%s%s%s.conf%s.dat"%(
                        Pi[0],Pi[1],Pi[2],K[0],K[1],K[2],Kg[0],Kg[1],Kg[2],conf_id)
            self.new_file[i]=(filen)

    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        for t_se in range(ct.Nt):
            for i,mom in enumerate(self.mom):
                Pi,K =  mom[0]
                Kg  =  mom[1]
                self.corr_th[i,0,(t_sk -t_se)%ct.Nt] += PiK_K_3P(perams_u_th[t_se,t_se],
                 perams_u_th[t_sk,t_se],perams_s_th[t_se],
                 VdV_th[VdV_idx[Pi],t_se],VdV_th[VdV_idx[K],t_se],
                            VdV_th[VdV_idx[Kg],t_sk],self.gamms)
class corr_fun_pikkg_mul(corr_fun_mul):
    def __init__(self,mom_pik,mom_kg,dire,gammas1,gammas2,
                 conf_id,VdV_th):
        corr_fun.__init__(self,mom_pik,dire)
        self.gamms1=gammas1
        self.gamms2=gammas2
        self.mom_pik=mom_pik
        self.mom_kg=mom_kg
        self.corr_th = th.zeros((len(self.gamms1), len(self.gamms2),
                                 len(self.mom_pik), ct.Nt), dtype=th.complex128,device=device)

        self.VdV_se1 = th.zeros((ct.Nt,ct.Nev1,ct.Nev1,len(self.mom_pik)), dtype=th.complex128,
                           device=device)
        self.VdV_se2 = th.zeros((ct.Nt,ct.Nev1,ct.Nev1,len(self.mom_pik)), dtype=th.complex128,
                           device=device)

        self.Gam1 = th.zeros((4,4,len(self.gamms1)),
                dtype=th.complex128, device=device)                   
        self.Gam2 = th.zeros((4,4,len(self.gamms2)), dtype=th.complex128,
                       device=device)                   
        for i , ga in enumerate(self.gamms1):
            self.Gam1[:,:,i] = Gammas[ga]
        for i , ga in enumerate(self.gamms2):
            self.Gam2[:,:,i] = Gammas[ga]
        for i,mompk in enumerate(self.mom_pik):
            Pi_se,K_se = mompk
            self.VdV_se1[:,:,:,i] = VdV_th[VdV_idx[Pi_se]]
            self.VdV_se2[:,:,:,i] = VdV_th[VdV_idx[K_se]]
            for j, ga1 in enumerate(self.gamms1):
                for l, ga2 in enumerate(self.gamms2):
                    filen= f"{self.dir}/corr_PiK_K"+ga1+ga2+\
                    "_Pi%s%s%sK%s%s%s_Kg%s%s%s.conf%s.dat"%( *Pi_se,*K_se,*self.mom_kg,conf_id)
                    self.new_file[(j,l,i)]=filen

    def contract(self,perams_u_th,perams_s_th,VdV_th,conf_id,t_sk):
        VdV_th_sk = VdV_th[VdV_idx[self.mom_kg]]
        for t_se in range(ct.Nt):
            self.corr_th[:,:,:,(t_sk -t_se)%ct.Nt] += PiK_K_3P_mul2(
                perams_u_th[t_se,t_se], perams_u_th[t_se],
                perams_s_th[t_sk,t_se], self.VdV_se1[t_se],self.VdV_se2[t_se],
                            VdV_th_sk[t_sk],self.Gam1,self.Gam2)
